#include "video_recorder.h"
#include "shared\dbg.h"

#include <qedit.h>
//#include <Streams.h>
#include "avi\avifile.h"

namespace avlib{
#define ERROR_IF_FAILED(str) \
 { if(FAILED(hr)) \
	{ \
	log(str);\
	return false;\
	}\
 }

#define RETURN_IF_FAILED(str) \
{ if(FAILED(hr)) \
	{ \
	log(str);\
	return hr;\
	}\
 }


VideoRecorder::VideoRecorder(){
	m_pBuilder = NULL;
	m_pGraph = NULL;

	m_width = 0;
	m_height = 0;
	m_pBits = NULL;
	m_timer = 0;

	m_avi = 0;
}

VideoRecorder::~VideoRecorder(){
	close();
	
}


bool VideoRecorder::init(){

	assert(m_pBuilder == NULL && m_pGraph == NULL);

	HRESULT hr = m_pBuilder.CoCreateInstance(CLSID_CaptureGraphBuilder2);
	ERROR_IF_FAILED("m_pBuilder.CoCreateInstance ");

	hr = m_pGraph.CoCreateInstance(CLSID_FilterGraph);
	ERROR_IF_FAILED(" m_pGraph.CoCreateInstance "); 

	hr = m_pBuilder->SetFiltergraph(m_pGraph);
	ERROR_IF_FAILED("m_pBuilder->SetFiltergraph "); 

	// Obtain interfaces for media control and Video Window
	hr = m_pGraph->QueryInterface(IID_IMediaControl, (LPVOID *) & m_pMediaControl);
	ERROR_IF_FAILED("m_pBuilder->SetFiltergraph "); 

	hr = m_pGraph->QueryInterface(IID_IVideoWindow, (LPVOID *) & m_pVideoWindow);
	ERROR_IF_FAILED("m_pBuilder->SetFiltergraph "); 

	hr = m_pGraph->QueryInterface(IID_IMediaEvent, (LPVOID *) & m_pMediaEvent);
	ERROR_IF_FAILED("m_pBuilder->SetFiltergraph "); 


	return true;
}

HRESULT VideoRecorder::GetCaptureFilter(int  idx,IBaseFilter ** filter){

	CComPtr<ICreateDevEnum> pDevEnum;
	CComPtr<IEnumMoniker> pEnumMoniker;
	CComPtr<IMoniker>	pMoniker;
	CComPtr<IBaseFilter>	pBase;
	//CComPtr<IPropertyBag>	pProperty;


	HRESULT hr = pDevEnum.CoCreateInstance(CLSID_SystemDeviceEnum);
	RETURN_IF_FAILED("pDevEnum.CoCreateInstance");

	hr = pDevEnum->CreateClassEnumerator(CLSID_VideoInputDeviceCategory, &pEnumMoniker, 0);
	RETURN_IF_FAILED("pDevEnum->CreateClassEnumerator");

	if(pEnumMoniker == NULL){
		log("no capture device");
		return E_FAIL;
	}

	for(int i=0;i<idx;i++){
		pMoniker = NULL;
		hr = pEnumMoniker->Next(1,&pMoniker,NULL);
		RETURN_IF_FAILED(" pEnumMoniker->Next");
	}

	if(pMoniker){
		hr = pMoniker->BindToObject(0, 0, IID_IBaseFilter, (void**)filter);
		RETURN_IF_FAILED("pMoniker->BindToObject");
	}

	return hr;
};



HRESULT VideoRecorder::SetupVideoWindow(HWND hWnd)
{
	HRESULT hr = S_OK;

	if(hWnd){
	// Set the window handle used to process graph events
		if(m_pMediaEvent){
			hr = m_pMediaEvent->SetNotifyWindow((OAHWND)hWnd, WM_GRAPHNOTIFY, 0);
			RETURN_IF_FAILED("m_pMediaEvent->SetNotifyWindow");
		}
		// Set the video window to be a child of the main window
		if(m_pVideoWindow){
			hr = m_pVideoWindow->put_Owner((OAHWND)hWnd);
			RETURN_IF_FAILED("m_pVideoWindow->put_Owner");
	
			hr = m_pVideoWindow->put_WindowStyle(WS_CHILD | WS_CLIPCHILDREN);
			RETURN_IF_FAILED("m_pVideoWindow->put_Owner");
			RECT rc;
			// Make the preview video fill our window
			GetClientRect(hWnd, &rc);
			m_pVideoWindow->SetWindowPosition(0, 0, rc.right, rc.bottom);
			// Make the video window visible, now that it is properly positioned

			hr = m_pVideoWindow->put_Visible(OATRUE);
			RETURN_IF_FAILED(" m_pVideoWindow->put_Visible");
		}
	}

	return hr;
}

void _FreeMediaType(AM_MEDIA_TYPE& mt)
{
	if (mt.cbFormat != 0)
	{
		CoTaskMemFree((PVOID)mt.pbFormat);
		mt.cbFormat = 0;
		mt.pbFormat = NULL;
	}
	if (mt.pUnk != NULL)
	{
		// pUnk should not be used.
		mt.pUnk->Release();
		mt.pUnk = NULL;
	}
}
void _DeleteMediaType(AM_MEDIA_TYPE *pmt)
{
	if (pmt != NULL)
	{
		_FreeMediaType(*pmt); 
		CoTaskMemFree(pmt);
	}
}


bool VideoRecorder::query_supported_dimension(int idx,std::set<DWORD> & dimensions)
{

	CComPtr<IBaseFilter> pBase;

	HRESULT hr = GetCaptureFilter(idx,&pBase);
	ERROR_IF_FAILED("GetCaptureFilter");	


	assert(pBase);

	CComPtr<IEnumPins> pEnum;
	CComPtr<IPin> pPin;
	CComPtr<IAMStreamConfig> pConfig;


	hr = pBase->EnumPins(&pEnum);
	//RETURN_IF_FAILED("enum pin");
	if(FAILED(hr)) return false;

	while(pEnum->Next(1,&pPin,0) == S_OK){

		hr = pPin->QueryInterface(IID_IAMStreamConfig,(void**)&pConfig);
		if(SUCCEEDED(hr)){
			int cnt,sz;

			hr = pConfig->GetNumberOfCapabilities(&cnt,&sz);
			AM_MEDIA_TYPE *p_mt;
			BYTE *pScc = new BYTE[sz];

			for(int i=0;i<cnt;i++){
				pConfig->GetStreamCaps(i,&p_mt,pScc);

				VIDEOINFOHEADER *pVih = reinterpret_cast<VIDEOINFOHEADER*>(p_mt->pbFormat);
				//log("support: %d * %d",pVih->bmiHeader.biWidth,pVih->bmiHeader.biHeight);
				
				dimensions.insert(MAKELONG((WORD)pVih->bmiHeader.biHeight,(WORD)pVih->bmiHeader.biWidth));
				_DeleteMediaType(p_mt);//in Strmbase.lib 
			}

			delete[] pScc;
			break;
		}
	}
	return hr==S_OK;
}

HRESULT VideoRecorder::EnumCaptureCapability(){

	HRESULT hr = S_OK;
	assert(m_pCapture);

	CComPtr<IEnumPins> pEnum;
	CComPtr<IPin> pPin;
	CComPtr<IAMStreamConfig> pConfig;


	hr = m_pCapture->EnumPins(&pEnum);
	RETURN_IF_FAILED("enum pin");

	while(pEnum->Next(1,&pPin,0) == S_OK){

		hr = pPin->QueryInterface(IID_IAMStreamConfig,(void**)&pConfig);
		if(SUCCEEDED(hr)){
			int cnt,sz;

			hr = pConfig->GetNumberOfCapabilities(&cnt,&sz);
			AM_MEDIA_TYPE *p_mt;
			BYTE *pScc = new BYTE[sz];

			for(int i=0;i<cnt;i++){
				pConfig->GetStreamCaps(i,&p_mt,pScc);

				VIDEOINFOHEADER *pVih = reinterpret_cast<VIDEOINFOHEADER*>(p_mt->pbFormat);
				log("support: %d * %d",pVih->bmiHeader.biWidth,pVih->bmiHeader.biHeight);

				_DeleteMediaType(p_mt);//in Strmbase.lib 
			}

			delete[] pScc;
			break;
		}
	}
	return hr;
}

HRESULT VideoRecorder::SetCaptureSize(IBaseFilter *pCapture){

	CComPtr<IAMStreamConfig > pConfig;

	HRESULT hr = m_pBuilder->FindInterface(&PIN_CATEGORY_CAPTURE,&MEDIATYPE_Video,pCapture, IID_IAMStreamConfig, (void**)&pConfig);
	RETURN_IF_FAILED("pBuilder->FindInterface");


	AM_MEDIA_TYPE *pCapmt;
	hr=pConfig->GetFormat(&pCapmt);
	RETURN_IF_FAILED("pBuilder->FindInterface");

	if(pCapmt->pbFormat){
		VIDEOINFOHEADER *pVih = reinterpret_cast<VIDEOINFOHEADER*>(pCapmt->pbFormat);
		pVih->bmiHeader.biWidth = m_width;
		pVih->bmiHeader.biHeight = m_height;
		int FRAMERATE=pVih->AvgTimePerFrame;


		hr=pConfig->SetFormat(pCapmt);
		RETURN_IF_FAILED("pConfig->SetFormat");
	}
	return hr;
}



bool VideoRecorder::open(int idx,HWND preview){
	//CLSID_VideoInputDeviceCategory

	m_wndPreview = preview;

	//CComPtr<IBaseFilter> pCapture;
	CComPtr<IBaseFilter> pGrabberFilter;
	//CComPtr<ISampleGrabber> pGrabber;

	HRESULT hr = GetCaptureFilter(idx,&m_pCapture);
	ERROR_IF_FAILED("GetCaptureFilter");	


	//EnumCaptureCapability();

	//capture filter can not create this way
	/*HRESULT hr = CoCreateInstance(clsid, NULL,CLSCTX_INPROC_SERVER,IID_PPV_ARGS(&pFilter));
	if(FAILED(hr)){
	log("CoCreateInstance error");
	return false;
	}*/

	hr = m_pGraph->AddFilter(m_pCapture,L"Capture device");
	ERROR_IF_FAILED("m_pGraph->AddFilter");	

	hr = SetCaptureSize(m_pCapture);
	ERROR_IF_FAILED("SetCaptureSize");	


	hr = pGrabberFilter.CoCreateInstance(CLSID_SampleGrabber);
	ERROR_IF_FAILED("pGrabberFilter.CoCreateInstance");	


	hr = m_pGraph->AddFilter(pGrabberFilter,L"Sample Grabber");
	ERROR_IF_FAILED("pGrabberFilter.CoCreateInstance");

	/*IBaseFilter *pNull = NULL;
	hr = CoCreateInstance(CLSID_NullRenderer, NULL, CLSCTX_INPROC_SERVER, IID_IBaseFilter, reinterpret_cast<void**>(&pNull));
	hr = mpGraph->AddFilter(pNull, L"NullRenderer")*/

	//set grabber format
	hr = pGrabberFilter->QueryInterface(IID_ISampleGrabber, (void**)&m_pGrabber);
	ERROR_IF_FAILED("pGrabberFilter.CoCreateInstance");
	AM_MEDIA_TYPE mt;
	ZeroMemory(&mt, sizeof(AM_MEDIA_TYPE));
	mt.majortype = MEDIATYPE_Video;
	//mt.subtype = MEDIASUBTYPE_RGB24;
	mt.subtype = MEDIASUBTYPE_RGB32;
	hr = m_pGrabber->SetMediaType(&mt);
	ERROR_IF_FAILED(" pGrabber->SetMediaType");

	hr = m_pGrabber->SetOneShot(FALSE);
	ERROR_IF_FAILED(" m_pGrabber->SetOneShot");

	hr = m_pGrabber->SetBufferSamples(TRUE);
	ERROR_IF_FAILED(" m_pGrabber->SetBufferSamples");


	hr = m_pBuilder->RenderStream(&PIN_CATEGORY_PREVIEW, &MEDIATYPE_Video, m_pCapture, pGrabberFilter, NULL);
	ERROR_IF_FAILED("  m_pBuilder->RenderStream");

	if(hr == 0x4027e){
		log("no preview pin,so a smartee added");
	}

	hr = m_pMediaControl->Run();
	ERROR_IF_FAILED("  m_pMediaControl->Run");

	OAFilterState pfs;
	while(true){
		hr = m_pMediaControl->GetState(100, &pfs);
		if( hr == S_OK) break;
		log("sleep for capture ready");
		Sleep(100);

	}
	SetupVideoWindow(m_wndPreview);
	return true;

}


bool VideoRecorder::init_bmp(int w,int h){

	m_width = w;
	m_height = h;

	buff_sz = w*h*4;

	m_pBits = malloc(w*h*4);
	log("alloc video buffer size %d\n",w*h*4);   
	memset(m_pBits,255,sizeof(w*h*4));

	
	return true;
}

void WINAPI onTimerVideoFunc(UINT wTimerID, UINT msg,DWORD dwUser,DWORD dwl,DWORD dw2){

	VideoRecorder *recorder = (VideoRecorder*)dwUser;

	recorder->blt();
}

void VideoRecorder::blt(){
	long sz=0;

	//EnterCriticalSection(&cs);
	HRESULT hr = m_pGrabber->GetCurrentBuffer(&sz,NULL);
	if(FAILED(hr)){
		log(" m_pGrabber->GetCurrentBuffer erro");
	}else{
		//assert(sz <= m_width*m_height*4);
		//sz=m_width*m_height*4;
		if(sz > buff_sz){
			log(" \n\nFATAL: video buffer size too small\n\n");
			free(m_pBits);
			log(" \n\nFATAL: original %d new %d \n\n",buff_sz,sz);
			m_pBits = malloc(sz);
			buff_sz = sz;
		}

		 hr = m_pGrabber->GetCurrentBuffer(&sz,(long*)m_pBits);
		if(FAILED(hr)){
			memset(m_pBits,0,m_width*m_height*4);
			log(" m_pGrabber->GetCurrentBuffer 2 erro");
		}

		if(m_avi)
			m_avi->AppendVideoFrame(m_pBits,0);
	}
	//LeaveCriticalSection(&cs);
}


bool VideoRecorder::start(AviFile *avi){

	m_avi = avi;
	
	//set format
	AM_MEDIA_TYPE mt;
	HRESULT hr = m_pGrabber->GetConnectedMediaType(&mt);
	ERROR_IF_FAILED(" m_pGrabber->GetConnectedMediaType");
	VIDEOINFOHEADER *pVih = (VIDEOINFOHEADER *)mt.pbFormat;
	int gChannels = pVih->bmiHeader.biBitCount / 8;
	int gWidth = pVih->bmiHeader.biWidth;
	int gHeight = pVih->bmiHeader.biHeight;
	_FreeMediaType(mt);

	//init bits
	init_bmp(gWidth,gHeight);
	
	//InitializeCriticalSection(&cs);
	//blt();
	//start timer
	m_timer = timeSetEvent(1000/m_fps,1000/m_fps,onTimerVideoFunc,(DWORD)this,TIME_PERIODIC|TIME_CALLBACK_FUNCTION|TIME_KILL_SYNCHRONOUS);
	if(m_timer == 0){
		log("start capture timer error %d",GetLastError());
		return false;
	}
	return true;
}

/*void *VideoRecorder::get_bits(){
	EnterCriticalSection(&cs);
	return m_pBits;
}

void VideoRecorder::release_bits(){
	LeaveCriticalSection(&cs);
}
*/

bool VideoRecorder::stop(){

	if(m_timer!= 0){
		
		//DeleteCriticalSection(&cs);
		timeKillEvent(m_timer);
		m_timer  = 0;
	}
	if(m_pMediaControl){
		HRESULT hr = m_pMediaControl->Pause();
		hr = m_pMediaControl->StopWhenReady();
		hr = m_pMediaControl->Stop();

		//wait for complete
		OAFilterState fs ;
		while (true){
			hr = m_pMediaControl->GetState(100, &fs);
			if (hr != VFW_S_STATE_INTERMEDIATE)
				break;
		}
	}

	if(m_pBits) {
		free(m_pBits);
		m_pBits = NULL;
	}
	
	SAFE_RELEASE(m_pBasicAudio);
	SAFE_RELEASE(m_pMediaSeeking);
	SAFE_RELEASE(m_pMediaEvent);
	SAFE_RELEASE(m_pMediaControl);
	SAFE_RELEASE(m_pVideoWindow);
	SAFE_RELEASE(m_pGrabber);
	SAFE_RELEASE(m_pCapture);
	SAFE_RELEASE(m_pGraph);
	SAFE_RELEASE(m_pBuilder);

	m_avi = 0;
	return true;
}

void VideoRecorder::close(){

	stop();
	
}

}